from ultralytics import YOLO
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
m = YOLO("wheel_seg_best.pt")
m.to(device)

success = m.export(format="onnx")  # export the model to ONNX format
